# -*- coding:utf-8 -*-#
# @Time:2023/5/31 14:48
# @Author:Adong
# @Software:PyCharm


"""
重要文件，谨慎修改！
重要文件，谨慎修改！
重要文件，谨慎修改！
重要文件，谨慎修改！
重要文件，谨慎修改！
"""


import os
import wave
import librosa
import numpy as np
import matplotlib.pyplot as plt
import random
from PIL import Image
from numpy import ndarray
from sklearn.preprocessing import MinMaxScaler
from scipy.fftpack import fft, ifft
from scipy.signal import savgol_filter
import scipy.io.wavfile as _wavefile


class Universal_tool:
    @classmethod
    def case_create(cls, target_folder, fault_type):
        """
        在目标文件夹target_folder创建各类故障fault_type的文件夹。
        :param fault_type: 期望创建文件夹的故障种类
        :param target_folder:[wav_to_gray,wav_to_MEL,wav_to_STFT]
        :return:
        """
        if not os.path.exists('./data/' + target_folder + '/'):  # 判断所在目录下是否有该文件名的文件夹
            os.mkdir('./data/' + target_folder + '/')  # 创建多级目录用mkdirs，单击目录mkdir
        set_type = ['train', 'test', 'veri']  # 设置三种类型的数据集
        for tt in set_type:  # 遍历set_type
            if not os.path.exists('./data/' + target_folder + '/' + tt):  # 判断所在目录下是否有该文件名的文件夹
                os.mkdir('./data/' + target_folder + '/' + tt)  # 创建多级目录用mkdirs，单击目录mkdir
            for t in fault_type:  # 遍历fault_type
                if not os.path.exists('./data/' + target_folder + '/' + tt + '/' + t):  # 判断所在目录下是否有该文件名的文件夹
                    os.mkdir('./data/' + target_folder + '/' + tt + '/' + t)  # 创建多级目录用mkdirs，单击目录mkdir

    '''FFT可以用来过滤低幅频率'''
    @classmethod
    def FFT_denoisying(cls, y: list, sr: int,if_show=False):
        """
        FFT降噪:
        FFT->转为对数频域->过滤下60%部分->IFFT
        :param y:原时域信号
        :param sr: 原信号采样频率
        :return: FFT降噪后的时域信号
        """
        '''原始时域信号绘图'''
        t = len(y) / sr  # 计算音频持续时长
        '''原始频域信号绘图'''
        y_fft = fft(y)  # 对y做傅里叶变换
        y_fft = y_fft * (2 / len(y))  # 幅值变换
        y_fft[0] = y_fft[0] / 2  # 最小频率直流信号的处理
        y_fft[len(y_fft) - 1] = y_fft[len(y_fft) - 1] / 2  # 最大频率直流信号的处理
        frequency = np.linspace(0, sr / 2 - 1, int(sr / 2))  # 奈奎斯特采样定理，有效频率范围为0~sr/2
        '''降噪频域信号绘图'''
        y_fft_real = np.abs(y_fft.copy())  # 读取y_fft的实部
        y_fft_real_log = np.log10(y_fft_real.copy())  # 取对数
        slide = min(y_fft_real_log)  # 计算平移长度
        y_fft_real_log = y_fft_real_log - slide  # 向y轴正方向平移，直到所有的值都非负
        indices = y_fft_real_log < np.max(y_fft_real_log.copy()) * 0.5  # 筛选对数实部值位于下60%的频率
        y_fft_denoisy = indices * y_fft.copy() * 0 + ~indices * y_fft.copy()  # 下60%幅值频率部分降低到0%，上40%幅值频率不变
        '''降噪时域信号绘图'''
        y_ifft = y_fft_denoisy.copy()  # 把降噪频域信号赋值给y_ifft
        y_ifft[0] = y_ifft[0] * 2  # 最小频率直流信号逆处理
        y_ifft[len(y_ifft) - 1] = y_ifft[len(y_ifft) - 1] * 2  # 最大频率直流信号逆处理
        y_ifft = y_ifft * (len(y) / 2)  # 幅值逆变换
        y_ifft = ifft(y_ifft)  # 逆傅里叶变换
        '''绘图'''
        if if_show is True:
            fig = plt.figure()
            ax1 = fig.add_subplot(231)  # 画原始时域信号
            ax2 = fig.add_subplot(232)  # 画原始频域信号
            ax3 = fig.add_subplot(233)  # 画原始对数频域
            ax4 = fig.add_subplot(234)  # 画降噪时域信号
            ax5 = fig.add_subplot(235)  # 画降噪频域信号
            ax6 = fig.add_subplot(236)  # 画降噪对数频域
            ax1.plot(np.linspace(0, t, len(y)), y)  # 横坐标时间，纵坐标原始时域幅值
            ax2.plot(frequency, np.abs(y_fft[0:int(sr / 2)]))  # 横坐标频率，纵坐标原始频域幅值
            ax3.plot(frequency, y_fft_real_log[0:int(sr / 2)])  # 横坐标频率，纵坐标原始对数频域
            ax4.plot(np.linspace(0, t, len(y_ifft)), np.real(y_ifft))  # 横坐标时间，纵坐标降噪时域幅值
            ax5.plot(frequency, np.abs(y_fft_denoisy[0:int(sr / 2)]))  # 横坐标频率，纵坐标降噪频域幅值
            ax6.plot(frequency, np.log10(np.abs(y_fft_denoisy[0:int(sr / 2)])) - slide)  # 横坐标频率，纵坐标降噪对数频域
            plt.show(block=True)
        return np.real(y_ifft)

    '''SMD可以用来去除瞬态干扰'''
    @classmethod
    def REPET(cls, y, sr):
        n_fft = int(25 / 1000 * sr)
        hop_length = int(10 / 1000 * sr)
        noise_stft = librosa.stft(y, n_fft=n_fft, hop_length=hop_length, window="hamm")  # stft得到时频谱
        '''相似矩阵滤波(论文法)'''
        from sklearn.metrics.pairwise import cosine_similarity
        import heapq
        noise_stft_trans = np.transpose(noise_stft.copy()).real # 将noise_stft转置并取实数值
        similar_matrix = cosine_similarity(noise_stft_trans)    # 计算相似矩阵
        num_of_maxs = int(0.8*len(similar_matrix))              # 设置相似帧的数量
        denoise_stft = []                                       # 存储更新后的时频图
        for idx1,row in enumerate(similar_matrix):              # 第idx1帧，row是第idx1帧的相似度向量
            row = row.tolist()                                  # array转list
            idx = list(map(row.index, heapq.nlargest(num_of_maxs,row)))  # 求最大的num_of_maxs个索引,nsmallest与nlargest相反，求最小
            similar_frame = noise_stft[:,idx]                  # 获得与取出帧相似度最大的若干帧
            V_sim = np.median(similar_frame, axis=1)           # np.real(noise_stft)每行的中位数。axis=1对应行，axis=0对应列
            temp = []                                          # 用于存储一帧的频率幅值
            for idx2,V in enumerate(V_sim):                    # 遍历每一个频率的中值
                if abs(noise_stft[idx2,idx1].real) <= abs(V.real):  # 如果第idx1帧的idx2频率的幅值小于中值，则保留原值
                    temp.append(noise_stft[idx2,idx1])
                else:
                    temp.append(V)
            denoise_stft.append(temp)
        denoise_stft = np.transpose(denoise_stft)
        '''计算维纳滤波系数部分'''
        # H_sim = denoise_stft / noise_stft  # 维纳滤波系数,H_sim是复数矩阵
        '''哈曼达积时频掩蔽'''
        # back_spec = noise_stft * H_sim  # 背景时频谱（振动信号），复数矩阵
        back_spec = denoise_stft
        noise_spec = noise_stft - back_spec  # 噪声信号，复数矩阵
        '''ISTFT部分'''
        istft_y = librosa.istft(noise_stft, n_fft=n_fft, hop_length=hop_length, window="hamm")
        istft_back_spec = librosa.istft(back_spec, n_fft=n_fft, hop_length=hop_length, window="hamm")
        istft_noise_spec = librosa.istft(noise_spec, n_fft=n_fft, hop_length=hop_length, window="hamm")
        # 含噪时频、背景时频、噪声时频、去噪时域、背景时域、噪声时域
        return noise_stft, back_spec, noise_spec, istft_y, istft_back_spec, istft_noise_spec

    '''savgol可以用来平滑曲线'''
    @classmethod
    def savgol(cls, y, sr,if_show = False):
        """savgol平滑部分"""
        time = len(y) / sr  # 计算音频时长
        t = np.linspace(0, time, len(y))
        y1 = savgol_filter(y, window_length=51, polyorder=2)
        '''画图部分'''
        if if_show is True:
            fig, ax = plt.subplots(figsize=(8, 4))
            ax.plot(t, y, label='Original Signal', lw=6, ls='--')
            ax.plot(t, y1, lw=3, label='Filtered Signal', color='C1')
            ax.set_xlabel('Time [s]')
            ax.set_ylabel('y')
            ax.legend(frameon=True, facecolor='white', fancybox=True)
            fig.tight_layout()
            plt.show(block=True)
        return y1

    @classmethod
    def array2wav(cls,nparray, framerate: int, wav_file: str):
        """
        将numpy数组转为单通道wav文件
        :param nparray: 输入的numpy向量
        :param wav_file: wav文件名
        :param framerate: 采样率(默认=sr)
        :return:
        """
        _wavefile.write(wav_file, framerate, nparray.astype(np.float32))

    @classmethod
    def AVF(cls,y,sr):
        n_fft = int(0.2 * sr)  # 帧长度是0.2s
        hop_length = int(0.02 * sr)    # 帧移0.02s
        noise_stft = librosa.stft(y,n_fft = n_fft,hop_length=hop_length,window="hamm")
        amp = noise_stft.real
        amp_var = []            # 所有频率幅值的方差
        for hz in amp:
            amp_var.append(np.var(hz))
        # plt.figure()
        # plt.plot(amp_var)
        # plt.show(block=True)

        '''计算50hz倍频点附近频点的方差值平均'''
        hz_50 = np.arange(50,len(noise_stft),50)        # 找出50hz倍频
        W_win = 10
        alphaX = []         # 保存50hz倍频附近W_win个频率点方差的平均值
        Ptrans = []             # 变压器频率
        Pfan = []               # 风扇频率
        for idx,x in enumerate(hz_50):
            piece = amp_var[x-W_win:x+W_win]        # 找出50hz倍频附近W_win个频率点的方差
            mean = np.mean(piece)/2                 # 计算平均值
            alphaX.append(mean)
            if mean > amp_var[x]:           # 如果平均值大于当前50倍频，则认为是本体
                Ptrans.append(x)
            else:
                Pfan.append(x)
        y_fft = fft(y)  # 对y做傅里叶变换
        y_fft_new = y_fft.copy()
        symmetry = len(y_fft)/2
        for idx,x in enumerate(y_fft):
            if idx in Pfan:
                y_fft_new[idx] = 0 + 0j
                y_fft_new[-idx] = 0 + 0j
            else:
                y_fft_new[idx] = x

        y_ifft = ifft(y_fft_new)  # 逆傅里叶变换
        plt.figure()
        ax1 = plt.subplot(221)
        ax2 = plt.subplot(222)
        ax3 = plt.subplot(223)
        ax4 = plt.subplot(224)
        ax1.plot(y_fft[:int(len(y_fft)/2)])
        ax2.plot(y_fft_new[:int(len(y_fft)/2)])
        ax3.plot(y)
        ax4.plot(y_ifft)
        plt.show(block=True)
        Universal_tool.array2wav(y_ifft, 'AVF.wav', sr)
        return y_ifft




class Transform2Gray:
    """
    wav文件转灰度图
    """

    def __init__(self, file_root):
        """
        :param file_root:wav文件的文件夹路径
        """
        self.pic_size = None
        self.file_root = file_root
        self.file_path_list, self.file_name_list = self.get_file_path_list()

    def get_file_path_list(self):
        """
        获取wav文件夹内所有的文件名file_name和文件路径file_path_list
        :return: 文件路径列表file_path_list,文件名列表file_name
        """
        file_path_list = []  # 存放wav文件路径的列表
        file_name = os.listdir(self.file_root)  # 获取wav文件的文件名
        for file in file_name:  # 遍历wav文件的文件名
            file_path = os.path.join(self.file_root, file)  # 把文件夹路径和文件名连接，组成文件路径
            file_path_list.append(file_path)  # 添加文件路径到列表
        return file_path_list, file_name

    def create_GRAY_1channel(self, filepath, pic_size, train_percent, test_percent, veri_percent, savepath,
                             if_FFT=False, if_savgol=False, if_SMD=False,if_AVF = False):
        """
        生成64*64的单通道灰度图
        :param if_SMD: 设置是否执行SMD降噪
        :param filepath:待转换文件的保存路径
        :param pic_size:转换图像的边长
        :param train_percent:划分训练集的比例
        :param test_percent:划分测试集的比例
        :param veri_percent:划分验证集的比例
        :param savepath:生成图片的保存文件夹路径
        :param if_FFT:设置是否执行FFT降噪
        :param if_savgol:设置是否执行savgol滤波
        :return:
        """
        '''音频片段获取部分'''
        y, sr = librosa.load(filepath)
        print("导入文件：" + filepath)
        sr = 48000
        y = librosa.resample(y=y, orig_sr=sr, target_sr=sr)[:sr]
        print("重采样完成!现采样率=" + str(sr))
        if if_SMD is True:
            data_stft, back_spec, noise_spec, istft_y, istft_back_spec, istft_noise_spec = Universal_tool.REPET(y, sr)
            y = istft_back_spec
            print("相似矩阵降噪完成!")
        if if_AVF is True:
            y = Universal_tool.AVF(y,sr)
            print("模值波动降噪完成!")
        if if_FFT is True:  # 设置是否应用FFT降噪
            y = Universal_tool.FFT_denoisying(y, sr)  # 对原信号执行FFT降噪操作
            print("FFT降噪完成!")
        if if_savgol is True:
            y = Universal_tool.savgol(y, sr)  # 对原信号执行savgol平滑操作
            print("savgol平滑完成!")

        d = librosa.get_duration(y=y, sr=sr)  # 获取音频时长
        test_length, veri_length = test_percent * d, veri_percent * d  # 设置用于测试集和验证集的音频长度
        train_length = train_percent * d  # 设置用于训练集的音频长度
        self.pic_size = pic_size  # 设置目标灰度图的边长
        # 获取训练集截取片段的起点,随机生成40个起点（也就是40张该类型的训练集图片）
        train_start = random.sample(range(0, round(train_length * sr - pic_size * pic_size)), 200)
        # 获取测试集截取片段的起点,随机生成5个起点
        test_start = random.sample(
            range(round(train_length * sr), round(train_length * sr + test_length * sr - pic_size * pic_size)), 25)
        # 获取验证集截取片段的起点,随机生成5个起点
        veri_start = random.sample(range(round((train_length + test_length) * sr), round(d * sr - pic_size * pic_size)),
                                   25)
        '''音频片段数据集生成'''
        train_dataset = []  # 训练集图片
        test_dataset = []  # 测试集图片
        veri_dataset = []  # 验证集图片
        for i in train_start:  # 遍历所有训练集起点
            train_dataset.append(y[i:i + pic_size * pic_size])  # 把训练集片段加入列表
        for j in test_start:  # 同上
            test_dataset.append(y[j:j + pic_size * pic_size])
        for k in veri_start:  # 同上
            veri_dataset.append(y[k:k + pic_size * pic_size])
        '''数据集生成模式遍历'''
        mode_list = ['train', 'test', 'veri']
        for mode in mode_list:
            if mode == 'train':
                dataset = train_dataset
            elif mode == 'test':
                dataset = test_dataset
            elif mode == 'veri':
                dataset = veri_dataset
            else:
                dataset = None
            '''转灰度图'''
            img_save_path = savepath + mode + '/' + filepath.split('/')[-2] + '/'  # 在savepath路径下的mode文件夹的filepath对应故障类型文件夹
            for idx, piece in enumerate(dataset):
                n = len(os.listdir(img_save_path))  # 记录文件夹中已经存在的图片数量，避免重复生成覆盖
                pic = np.zeros(shape=(pic_size, pic_size))  # 生成空白灰度图
                for x in range(pic_size):  # 循环size次，遍历每一行的行号
                    rowstrat = pic_size * x  # 获取该行起点在上述音频片段中的索引
                    rowend = pic_size * (x + 1)  # 获取改行终点在上述音频片段中的索引
                    row = piece[rowstrat:rowend:1]  # 在上述音频片段中截取该行片段
                    pic[x] = row  # 将截取的该行片段放入空白灰度图中
                normalize_tool = MinMaxScaler(feature_range=(0, 255))  # 导入最大最小归一化工具，设置归一化范围为(0,255)
                normalized_pic = normalize_tool.fit_transform(pic)  # 执行归一化
                img = Image.fromarray(normalized_pic)  # 将array转为image
                img = img.convert('L')  # 将image转为灰度图
                img.save(img_save_path + str(n + 1) + '.png')  # 保存灰度图


class Transform2RGB:
    """
    wav文件转RGB
    """

    def __init__(self, file_root=None):
        self.filepath = file_root  # 音频文件地址
        self.y, self.sr = librosa.load(file_root)  # 读取音频的信号值和采样率
        self.audio_time = librosa.get_duration(path=file_root)  # 读取音频持续的时间
        the_length_of_the_audio_clip = 2.2  # 音频片段的时长
        self.samples = self.y[0:int(the_length_of_the_audio_clip / self.audio_time * len(self.y))]

    def time_domain(self):
        """
        生成音频时域图
        :return:
        """
        plt.figure()  # 生成画板
        ax1 = plt.subplot(1, 2, 1)  # 分区1：用于展示时域图
        ax2 = plt.subplot(1, 2, 2)  # 分区2：用于展示时域图转成的灰度图
        '''制作时域图'''
        time = np.linspace(0, self.audio_time, len(self.samples))  # 在0~self.audio_time的范围内生成len(self.samples)个点
        samples = self.samples
        ax1.plot(time, samples)
        ax1.set_title("time-amplitude")
        ax1.set_xlabel("time(s)")
        ax1.set_ylabel("amplitude(db)")
        '''制作灰度图'''
        m = min(samples)
        while len(samples) < 224 * 224:
            samples = np.append(samples, m)
        reshape_samples = samples[:224 * 224].reshape((224, 224))
        normalize_tool = MinMaxScaler(feature_range=(0, 255))  # 导入最大最小归一化工具，设置归一化范围为(0,255)
        normalized_pic = normalize_tool.fit_transform(reshape_samples)  # 执行归一化
        img = Image.fromarray(normalized_pic)  # 将array转为image
        img = img.convert('L')  # 将image转为灰度图
        ax2.imshow(img, cmap='gray')  # 显示单通道图像
        '''展示图像'''
        plt.show(block=True)
        return normalized_pic

    def frequency_domain(self):
        """
        生成音频频谱图
        :return:
        """
        plt.figure()  # 生成画板
        ax1 = plt.subplot(1, 2, 1)  # 分区1：用于展示时域图
        ax2 = plt.subplot(1, 2, 2)  # 分区2：用于展示时域图转成的灰度图
        '''制作频域图'''
        y_f = fft(self.samples)  # 快速傅里叶变换
        y_f = y_f * (2 / len(self.samples))  # 幅值变换
        y_f[0] = y_f[0] / 2  # 最小频率直流信号的处理
        y_f[len(y_f) - 1] = y_f[len(y_f) - 1] / 2  # 最大频率直流信号的处理
        frequency = np.linspace(0, self.sr / 2 - 1, int(self.sr / 2))  # 奈奎斯特采样定理，有效频率范围为0~sr/2
        ax1.plot(frequency, abs(y_f[0:int(self.sr / 2)]))
        ax1.set_title("frequency-amplitude")
        ax1.set_xlabel("frequency(Hz)")
        ax1.set_ylabel("amplitude(db)")
        '''制作灰度图'''
        y_f = abs(y_f)
        m = min(y_f)
        while len(y_f) < 224 * 224:
            y_f = np.append(y_f, m)
        reshape_magnitude = y_f[:224 * 224].reshape((224, 224))  # 取magnititude前224*224个点转换为224*224的矩阵，用与制作单通道图像
        normalize_tool = MinMaxScaler(feature_range=(0, 255))  # 导入最大最小归一化工具，设置归一化范围为(0,255)
        normalized_pic = normalize_tool.fit_transform(reshape_magnitude)  # 执行归一化
        img = Image.fromarray(normalized_pic)  # 将array转为image
        img = img.convert('L')  # 将image转为灰度图
        ax2.imshow(img, cmap='gray')  # 显示单通道图像
        # img.save('1.png')  # 保存灰度图
        '''展示图像'''
        plt.show(block=True)
        return normalized_pic

    def time_frequency_domain(self):
        """
        生成音频时频图
        :return:
        """
        fig = plt.figure()  # 生成画板
        ax1 = plt.subplot(1, 2, 1)  # 分区1：用于展示时域图
        ax2 = plt.subplot(1, 2, 2)  # 分区2：用于展示时域图转成的灰度图
        '''制作时频域图'''
        n_fft = int(self.sr)
        stft = librosa.stft(y=self.samples, n_fft=n_fft)
        y_f = stft * (2 / len(self.samples))  # 幅值变换
        y_f[:][0] = y_f[:][0] / 2  # 最小频率直流信号的处理
        y_f[:][len(y_f) - 1] = y_f[:][len(y_f) - 1] / 2  # 最大频率直流信号的处理
        '''FFT滤波'''
        # indices = stft > 300  # filter out those value under 300
        # stft = indices * stft  # noise frequency will be set to 0
        # spectrogram = librosa.amplitude_to_db(np.abs(stft))
        spectrogram = abs(stft) * (2 / len(self.samples))

        frequency = np.linspace(0, len(spectrogram) - 1, len(spectrogram))
        ax1.plot(frequency, spectrogram)
        ax1.set_title("frequency-amplitude")
        ax1.set_xlabel("frequency(Hz)")
        ax1.set_ylabel("amplitude(db)")
        '''制作灰度图'''
        col = spectrogram.shape[1]
        cut = round(224 * 224 / col)
        spectrogram = spectrogram[:cut]
        m = np.min(spectrogram)
        while len(spectrogram) < 224 * 224:
            spectrogram = np.append(spectrogram, m)
        reshape_spectrogram = spectrogram[:224 * 224].reshape(
            (224, 224))  # 取magnititude前224*224个点转换为224*224的矩阵，用与制作单通道图像
        normalize_tool = MinMaxScaler(feature_range=(0, 255))  # 导入最大最小归一化工具，设置归一化范围为(0,255)
        normalized_pic = normalize_tool.fit_transform(reshape_spectrogram)  # 执行归一化
        img = Image.fromarray(normalized_pic)  # 将array转为image
        img = img.convert('L')  # 将image转为灰度图
        ax2.imshow(img, cmap='gray')  # 显示单通道图像
        '''展示图像'''
        plt.show(block=True)
        return normalized_pic

    def array2RGB(self, red, green, blue):
        r = Image.fromarray(red).convert('L')
        g = Image.fromarray(green).convert('L')
        b = Image.fromarray(blue).convert('L')
        img = Image.merge('RGB', (r, g, b))
        img.show()
        img.save('1.png')


class Transform2Mel:
    '''
    把音频时域信号转成符合resnet输入尺寸的1*3*224*224的mel语谱图
    '''
    def __init__(self, file_root):
        """
        :param file_root:wav文件的文件夹路径
        """
        self.pic_size = None
        self.file_root = file_root
        self.file_path_list, self.file_name_list = self.get_file_path_list()

    def get_file_path_list(self):
        """
        获取wav文件夹内所有的文件名file_name和文件路径file_path_list
        :return: 文件路径列表file_path_list,文件名列表file_name
        """
        file_path_list = []  # 存放wav文件路径的列表
        file_name = os.listdir(self.file_root)  # 获取wav文件的文件名
        for file in file_name:  # 遍历wav文件的文件名
            file_path = os.path.join(self.file_root, file)  # 把文件夹路径和文件名连接，组成文件路径
            file_path_list.append(file_path)  # 添加文件路径到列表
        return file_path_list, file_name


    def basicData(self,y,sr):
        '''REPET盲源分离'''
        # noise_stft, back_spec, noise_spec, istft_y, istft_back_spec, istft_noise_spec = Universal_tool.REPET(y, sr)
        self.y = y
        self.sr = sr
        self.pic_size = 224
        self.n_fft = (self.pic_size - 1) * 2
        self.hop_length = self.pic_size - 1
        self.audio_length = (self.pic_size-1)*self.hop_length

    def cut_pic(self,train_num,test_num,veri_num):
        """
        窗长为audio_length，窗移为
        :param train_num:
        :param test_num:
        :param veri_num:
        :return:
        """
        audio_length = self.audio_length
        y = self.y
        train_pic = []
        test_pic = []
        veri_pic = []
        pics_begin = random.sample(range(0, round(len(y) - audio_length)), train_num + test_num + veri_num)
        for i in range(train_num):
            train_pic.append(y[pics_begin[i]:pics_begin[i]+audio_length])
        for j in range(train_num,train_num + test_num):
            test_pic.append(y[pics_begin[j]:pics_begin[j] + audio_length])
        for k in range(train_num+test_num,train_num + test_num + veri_num):
            veri_pic.append(y[pics_begin[k]:pics_begin[k] + audio_length])

        self.train_pic = train_pic
        self.test_pic = test_pic
        self.veri_pic = veri_pic
        return train_pic,test_pic,veri_pic
    def zero_padding(self,mel_spect):
        """
        把128*224的Mel图居中零填充到224*224的尺寸
        """
        mel_spect = mel_spect.tolist()
        zero = np.zeros(len(mel_spect[0])).tolist()
        need_to_pad = self.pic_size - len(mel_spect)
        front = int(need_to_pad / 2)
        back = need_to_pad - front
        for i in range(front):
            mel_spect.insert(0,zero)
        for j in range(back):
            mel_spect.append(zero)
        return np.array(mel_spect)

    def trans2singleMel(self,y):
        '''
        生成单张Mel语谱图
        :return:
        '''
        mel_spect = librosa.feature.melspectrogram(y=y, sr=self.sr, n_fft=self.n_fft, hop_length=self.hop_length)   # Mel谱
        noise_stft = librosa.stft(y,n_fft=self.n_fft, hop_length=self.hop_length, window="hamm").real  # stft得到时频谱


        normalize_tool = MinMaxScaler(feature_range=(0, 255))       # 导入最大最小归一化工具，设置归一化范围为(0,255)
        normalized_mel = normalize_tool.fit_transform(mel_spect)    # 执行归一化
        normalized_stft = normalize_tool.fit_transform(noise_stft)
        mel_spect = self.zero_padding(normalized_mel)
        noise_stft = self.zero_padding(normalized_stft)
        # img = Image.fromarray(normalized_pic)                       # 将array转为image
        # img = img.convert('L')                                      # 将image转为灰度图
        # img.save('zlpc.png')                                        # 保存灰度图
        return mel_spect,noise_stft
    def merge(self,pics,img_save_path):
        """
        随机取三个pic合并成RGB
        :param pics:[train_pic,test_pic,veri_pic]
        :param num:生成RGB图的数量
        :return:
        """
        for pic in pics:
            n = len(os.listdir(img_save_path))  # 记录文件夹中已经存在的图片数量，避免重复生成覆盖

            '''常规时频图、Mel时频图提取'''
            # mel_spect,noise_stft = self.trans2singleMel(pic)
            '''时域图提取'''
            need_to_pad = self.pic_size * self.pic_size - len(pic)
            front = int(need_to_pad / 2)
            back = need_to_pad - front
            time_array = np.pad(pic, (front, back), 'constant', constant_values=(0, 0))
            time_array = time_array.reshape(self.pic_size,self.pic_size)
            normalize_tool = MinMaxScaler(feature_range=(0, 255))  # 导入最大最小归一化工具，设置归一化范围为(0,255)
            time_array = normalize_tool.fit_transform(time_array)  # 执行归一化
            '''合成RGB'''
            # r = Image.fromarray(mel_spect).convert('L')
            # g = Image.fromarray(noise_stft).convert('L')
            b = Image.fromarray(time_array).convert('L')
            b.save(img_save_path + str(n) + '.png')
            # img = Image.merge('RGB', (r, g, b))
            # img.save(img_save_path + str(n) + '.png')

        '''三个不同时段的Mel图合成RGB'''
        # for i in range(len(pics)):
        #     n = len(os.listdir(img_save_path))  # 记录文件夹中已经存在的图片数量，避免重复生成覆盖
        #
        #     pic3 = random.sample(pics,3)
        #
        #     r = Image.fromarray(self.trans2singleMel(pic3[0])).convert('L')
        #     g = Image.fromarray(self.trans2singleMel(pic3[1])).convert('L')
        #     b = Image.fromarray(self.trans2singleMel(pic3[2])).convert('L')
        #     img = Image.merge('RGB', (r, g, b))
        #     # img.show()
        #     img.save(img_save_path + str(n) + '.png')


if __name__ == '__main__':
    '''生成灰度图数据集部分'''
    # dataset_verision = 'V4'  # 设置选择的数据集版本
    # type_list = os.listdir('./data/wav_data_' + dataset_verision)  # 读取该wav数据集下的故障类型
    # Universal_tool.case_create('wav' + dataset_verision + '_to_gray', type_list)  # 创建文件夹
    # for t in type_list:
    #     file_root = r'./data/wav_data_' + dataset_verision + '/' + t + '/'  # 生成图片用到的故障类型wav文件夹
    #     test = Transform2Gray(file_root)
    #     for item in test.file_path_list:
    #         test.create_GRAY_1channel(item, pic_size=64, test_percent=0.1, veri_percent=0.1, train_percent=0.8,
    #                                   savepath='./data/wav' + dataset_verision + '_to_gray/', if_FFT=False,
    #                                   if_savgol=False, if_SMD=False,if_AVF=False)

    '''生成RGB Mel图数据集部分'''
    dataset_verision = 'V3_test'  # 设置选择的数据集版本
    print("选择的数据集为：wav_data_" + dataset_verision)
    type_list = os.listdir('./data/wav_data_' + dataset_verision)  # 读取该wav数据集下的故障类型
    print("该数据集中的故障类型有：" + str(type_list))

    # Universal_tool.case_create('wav' + dataset_verision + '_to_RGB', type_list)  # 创建文件夹
    # print("已完成wav" + dataset_verision + "_to_Mel文件夹的创建")
    Universal_tool.case_create('wavV3_to_gray', type_list)  # 创建文件夹
    print("已完成wav" + dataset_verision + "_to_gray文件夹的创建")

    for t in type_list:
        print("正在生成" + str(t) + "类型")
        file_root = r'./data/wav_data_' + dataset_verision + '/' + t + '/'  # 生成图片用到的故障类型wav文件夹
        zqd = Transform2Mel(file_root)
        for item in zqd.file_path_list:
            y,sr = librosa.load(item)
            print(str(item) + "已加载")
            zqd.basicData(y, sr)
            zqd.cut_pic(100, 0, 0)
            zqd.merge(zqd.train_pic, img_save_path='./data/wavV3_to_gray/test/' + t + "/")
            # zqd.merge(zqd.train_pic, img_save_path='./data/wav' + dataset_verision + '_to_RGB/train/' + t + "/")
            # zqd.merge(zqd.test_pic, img_save_path='./data/wav' + dataset_verision + '_to_RGB/test/' + t + "/")
            # zqd.merge(zqd.veri_pic, img_save_path='./data/wav' + dataset_verision + '_to_RGB/veri/' + t + "/")
